from time import time
from pathlib import Path
import numpy as np
import cv2


from cardiovascular_monitoring.ugent.heartbeat_classification.data.mit_bih import DATA_ROOT, OUTPUT_FOLDER, LABEL_MAPPING, TRAINING_IDS, TESTING_IDS
from cardiovascular_monitoring.ugent.heartbeat_classification.data.mit_bih import read_record, remove_baseline_wander, calculate_fft, get_rr_data, calculate_cwt


def load_datasets():
    X_train = np.load(OUTPUT_FOLDER / 'rnn' / 'X_train.npy')
    X_fft_train = np.load(OUTPUT_FOLDER / 'rnn' / 'X_fft_train.npy')
    rr_train = np.load(OUTPUT_FOLDER / 'rnn' / 'rr_train.npy')
    y_train = np.load(OUTPUT_FOLDER / 'rnn' / 'y_train.npy')

    X_test = np.load(OUTPUT_FOLDER / 'rnn' / 'X_test.npy')
    X_fft_test = np.load(OUTPUT_FOLDER / 'rnn' / 'X_fft_test.npy')
    rr_test = np.load(OUTPUT_FOLDER / 'rnn' / 'rr_test.npy')
    y_test = np.load(OUTPUT_FOLDER / 'rnn' / 'y_test.npy')

    return X_train, X_fft_train, rr_train, y_train, \
           X_test, X_fft_test, rr_test, y_test


def get_windows(X, X_cwt, segments, n_before, n_after):
    X_windows = []
    X_cwt_windows = []
    X_fft_windows = []
    rr_windows = []
    y_windows = []

    for loc, y, rr in segments:
        if loc[0] < n_before or loc[-1] + n_after >= len(X):
            continue
        X_windows.append(np.asarray([X[l - n_before:l + n_after] for l in loc]))
        X_cwt_windows.append(np.asarray([cv2.resize(X_cwt[l - n_before:l + n_after], (100, 100)) for l in loc]))
        X_fft_windows.append(calculate_fft(X_windows[-1]))
        rr_windows.append(np.asarray(rr))
        y_windows.append(np.asarray([LABEL_MAPPING[l] for l in y]))

    return np.stack(X_windows), np.stack(X_cwt_windows), np.stack(X_fft_windows), np.stack(rr_windows), np.stack(y_windows)


def get_rr_data(loc):
    avg_rr = np.mean(np.diff(loc))
    rr_data = []
    for i in range(len(loc)):
        if i < 10 or i > len(loc) - 2:
            rr_data.append(None)
        else:
            prev_rr = loc[i] - loc[i - 1]
            next_rr = loc[i + 1] - loc[i]
            ratio_rr = prev_rr / next_rr
            local_rr = np.mean(np.diff(loc[i - 10:i + 1]))
            rr_data.append((prev_rr / avg_rr, next_rr / avg_rr, ratio_rr, local_rr / avg_rr))

    return rr_data


def process_dataset(data_root, patient_ids,
                    n_windows=10, lead='MLII', sampling_rate=360, n_before=90, n_after=110):
    X_total = []
    X_cwt_total = []
    X_fft_total = []
    rr_total = []
    y_total = []
    for i, record_id in enumerate(patient_ids):
        print(f"Processing patient {i + 1} / {len(patient_ids)}")
        X, y, loc = read_record(data_root, record_id, lead)
        X = remove_baseline_wander(X, sampling_rate)
        X = X / np.mean(loc)
        X_cwt = calculate_cwt(X, sampling_rate)
        rr_data = get_rr_data(loc)
        mask = [x is not None for x in rr_data]
        y, loc, rr_data = y[mask], loc[mask], [x for x, m in zip(rr_data, mask) if m]
        segments = [(loc[i:i + n_windows], y[i:i + n_windows], rr_data[i:i + n_windows]) for i in range(0, len(loc) - n_windows, 1)]

        X_windows, X_cwt_windows, X_fft_windows, rr_windows, y_windows = get_windows(X, X_cwt, segments,
                                                                                     n_before, n_after)
        X_total.append(X_windows)
        X_cwt_total.append(X_cwt_windows)
        X_fft_total.append(X_fft_windows)
        rr_total.append(rr_windows)
        y_total.append(y_windows)

    return np.concatenate(X_total), np.concatenate(X_cwt_total), np.concatenate(X_fft_total), np.concatenate(rr_total), np.concatenate(y_total)


if __name__ == '__main__':
    t_start = time()
    X_train, X_cwt_train, X_fft_train, rr_train, y_train = process_dataset(DATA_ROOT, TRAINING_IDS)
    t_end = time()
    print(X_train.shape, X_cwt_train.shape, X_fft_train.shape, rr_train.shape, y_train.shape)
    print(f"Processing training data took {t_end - t_start:.02f}s\n")

    t_start = time()
    X_test, X_cwt_test, X_fft_test, rr_test, y_test = process_dataset(DATA_ROOT, TESTING_IDS)
    t_end = time()
    print(X_test.shape, X_cwt_test.shape, X_fft_test.shape, rr_test.shape, y_test.shape)
    print(f"Processing testing data took {t_end - t_start:.02f}s\n")

    np.save(OUTPUT_FOLDER / 'rnn/X_train.npy', X_train)
    np.save(OUTPUT_FOLDER / 'rnn/X_cwt_train.npy', X_cwt_train)
    np.save(OUTPUT_FOLDER / 'rnn/X_fft_train.npy', X_fft_train)
    np.save(OUTPUT_FOLDER / 'rnn/rr_train.npy', rr_train)
    np.save(OUTPUT_FOLDER / 'rnn/y_train.npy', y_train)
    
    np.save(OUTPUT_FOLDER / 'rnn/X_test.npy', X_test)
    np.save(OUTPUT_FOLDER / 'rnn/X_cwt_test.npy', X_cwt_test)
    np.save(OUTPUT_FOLDER / 'rnn/X_fft_test.npy', X_fft_test)
    np.save(OUTPUT_FOLDER / 'rnn/rr_test.npy', rr_test)
    np.save(OUTPUT_FOLDER / 'rnn/y_test.npy', y_test)
